import os
from typing import Dict, List, Optional, TypedDict
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langgraph.graph import StateGraph, END
from tqdm import tqdm  # 进度条支持
from datetime import datetime
from dotenv import load_dotenv
import logging
import shutil
import platform
import time

# 根据操作系统选择锁机制
try:
    import portalocker

    def acquire_lock(lock_file):
        try:
            lock = portalocker.Lock(lock_file, timeout=2, fail_when_locked=True)
            return lock
        except portalocker.AlreadyLocked:
            logger.error("检测到未释放的锁，正在自动清理...")
            os.remove(lock_file)
            return portalocker.Lock(lock_file)
except ImportError:
    import fcntl
    def acquire_lock(lock_file):
        lock = open(lock_file, 'w')
        fcntl.flock(lock, fcntl.LOCK_EX)
        return lock

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 环境变量加载
load_dotenv()

# 定义一个配置类，用于管理各种参数
class RAGConfig:
    """可配置参数"""
    def __init__(self):
        self.model_path = os.getenv("MODEL_PATH", "D:/ideaSpace/MyPython/models/bge-small-zh-v1.5")
        self.chroma_dir = os.getenv("CHROMA_DIR", "D:/tmp/chroma_db")
        self.chunk_size = int(os.getenv("CHUNK_SIZE", 600))
        self.chunk_overlap = int(os.getenv("CHUNK_OVERLAP", 80))
        self.separators = ["\n\n", "\n", "。", "！", "？"]
        self.search_kwargs = {"k": 3}  # 检索返回数量
        self.force_rebuild = False  # 是否强制重建向量库

# 定义一个状态类，用于在工作流中传递状态
class RAGState(TypedDict):
    """工作流状态定义"""
    pdf_path: str
    documents: Optional[List[Dict]]
    vectorstore: Optional[Chroma]
    question: Optional[str]
    answer: Optional[str]
    context: Optional[List[str]]

class PDFRAGWorkflow:
    def __init__(self, config: Optional[RAGConfig] = None):
        self.config = config if config else RAGConfig()
        self.lock_file = os.path.join(self.config.chroma_dir, "rag.lock")
        self.llm, self.embeddings = self._init_models()
        self.workflow = self._build_workflow()

    def _init_models(self):
        """初始化模型实例（初始化语言模型和嵌入模型）"""
        return (
            ChatZhipuAI(
                api_key=os.getenv("ZHIPUAI_API_KEY"),
                model="glm-4",
                temperature=0.3
            ),
            HuggingFaceEmbeddings(
                model_name=self.config.model_path,
                model_kwargs={'device': 'cpu'}
            )
        )

    def _load_documents(self, state: RAGState) -> RAGState:
        """加载PDF文档（带进度显示）"""
        logger.info(f"加载PDF: {state['pdf_path']}")
        loader = PyPDFLoader(state["pdf_path"])
        pages = loader.load()

        # 显示加载进度
        state["documents"] = []
        for page in tqdm(pages, desc="解析PDF页面"):
            page.metadata.update({
                "source": state["pdf_path"],
                "load_time": datetime.now().isoformat()
            })
            state["documents"].append(page)
        return state

    def _split_documents(self, state: RAGState) -> RAGState:
        """智能文本分割"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.config.chunk_size,
            chunk_overlap=self.config.chunk_overlap,
            separators=self.config.separators
        )
        state["documents"] = splitter.split_documents(state["documents"])
        return state

    def _create_vectorstore(self, state: RAGState) -> RAGState:
        """创建/加载向量数据库"""
        try:
            lock = acquire_lock(self.lock_file)
            try:
                if os.path.exists(self.config.chroma_dir) and not self.config.force_rebuild:
                    logger.info("加载现有向量库")
                    state["vectorstore"] = Chroma(
                        persist_directory=self.config.chroma_dir,
                        embedding_function=self.embeddings
                    )
                else:
                    logger.info("创建新向量库")
                    if os.path.exists(self.config.chroma_dir):
                        self._force_remove_dir(self.config.chroma_dir)

                    state["vectorstore"] = Chroma.from_documents(
                        documents=state["documents"],
                        embedding=self.embeddings,
                        persist_directory=self.config.chroma_dir
                    )
            finally:
                if isinstance(lock, portalocker.Lock):
                    lock.release()
                else:
                    fcntl.flock(lock, fcntl.LOCK_UN)
                    lock.close()
        except Exception as e:
            logger.error(f"向量库操作失败: {str(e)}")
            raise
        return state

    def _force_remove_dir(self, dir_path):
        """强制删除目录（跨平台）"""
        max_retries = 3
        retries = 0
        while retries < max_retries:
            try:
                shutil.rmtree(dir_path)
                break
            except PermissionError:
                retries += 1
                logger.warning(f"删除目录失败，重试中 ({retries}/{max_retries})")
                time.sleep(1)
        else:
            logger.error(f"无法删除目录: {dir_path}")
            raise RuntimeError(f"无法删除目录: {dir_path}")

    def _safe_remove_dir(self, dir_path):
        """安全删除目录（自动处理权限和锁）"""
        max_retries = 3
        for _ in range(max_retries):
            try:
                if os.path.exists(dir_path):
                    # 1. 解除文件只读属性（Windows）
                    if platform.system() == "Windows":
                        os.system(f'attrib -r "{dir_path}\\*" /s')

                    # 2. 递归修改权限
                    for root, dirs, files in os.walk(dir_path):
                        for name in files + dirs:
                            path = os.path.join(root, name)
                            try:
                                os.chmod(path, 0o777)
                            except:
                                pass

                    # 3. 实际删除
                    shutil.rmtree(dir_path)
                    return True
            except Exception as e:
                logger.warning(f"删除重试中... 错误: {str(e)}")
                time.sleep(1)

        raise RuntimeError(
            f"无法删除目录 {dir_path}\n"
            f"请手动删除或检查权限:\n"
            f"终端执行: rm -rf {os.path.abspath(dir_path)}"
        )

    def _retrieve_context(self, state: RAGState) -> RAGState:
        """检索相关上下文"""
        retriever = state["vectorstore"].as_retriever(
            search_kwargs=self.config.search_kwargs
        )
        state["context"] = [
            f"来源：{doc.metadata.get('source', '未知')}\n"
            f"内容：{doc.page_content}"
            for doc in retriever.invoke(state["question"])
        ]
        return state

    def _generate_answer(self, state: RAGState) -> RAGState:
        """生成自然语言回答"""
        prompt_template = """
        根据以下上下文回答问题：
        {context}

        问题：{question}
        要求：
        1. 回答需简洁准确
        2. 标明引用来源
        3. 不确定时说明"根据现有信息无法确定"

        回答："""

        prompt = ChatPromptTemplate.from_template(prompt_template)
        chain = (
                {"context": RunnablePassthrough(),
                 "question": RunnablePassthrough()}
                | prompt
                | self.llm
                | StrOutputParser()
        )

        state["answer"] = chain.invoke({
            "context": "\n\n".join(state["context"]),
            "question": state["question"]
        })
        return state

    def _build_workflow(self):
        """构建LangGraph工作流"""
        workflow = StateGraph(RAGState)

        # 定义节点
        nodes = [
            ("load_docs", self._load_documents),
            ("split_docs", self._split_documents),
            ("create_vs", self._create_vectorstore),
            ("retrieve", self._retrieve_context),
            ("generate", self._generate_answer)
        ]
        for name, func in nodes:
            workflow.add_node(name, func)

        # 定义边
        workflow.add_edge("load_docs", "split_docs")
        workflow.add_edge("split_docs", "create_vs")

        # 条件分支
        workflow.add_conditional_edges(
            "create_vs",
            lambda state: "question" in state and bool(state["question"]),
            {True: "retrieve", False: END}
        )
        workflow.add_edge("retrieve", "generate")
        workflow.add_edge("generate", END)
        # 设置工作流程的起始节点
        workflow.set_entry_point("load_docs")
        return workflow.compile()

    def initialize(self, pdf_path: str, force_rebuild=False):
        """初始化文档处理"""
        self.config.force_rebuild = force_rebuild
        return self.workflow.invoke({"pdf_path": pdf_path})

    def query(self, question: str, state: RAGState) -> str:
        """执行查询"""
        state["question"] = question
        result = self.workflow.invoke(state)
        return result["answer"]

    def cleanup(self):
        """资源清理"""
        if hasattr(self, 'vectorstore'):
            self.vectorstore.delete_collection()
        if os.path.exists(self.lock_file):
            os.remove(self.lock_file)

# 使用示例
if __name__ == "__main__":
    rag = None
    try:
        # 初始化配置
        config = RAGConfig()
        config.chunk_size = 800  # 可自定义参数

        #创建实例
        rag = PDFRAGWorkflow(config)

        # 处理文档（首次运行force_rebuild=True，添加该变量后无法删除chroma_dir）
        state = rag.initialize("D:/document/阿里巴巴Java开发手册-嵩山版.pdf")

        # 交互问答
        print("系统初始化完成，输入问题开始查询（输入q退出）")
        while True:
            question = input("\n问题：").strip()
            if question.lower() == 'q':
                break

            answer = rag.query(question, state)
            print(f"\n回答：{answer}")

    except Exception as e:
        logger.exception(f"系统错误: {str(e)}")
    finally:
        if rag:
            rag.cleanup()